package com.gitee.gateway.core.utils;

import com.fasterxml.jackson.core.JsonFactory;
import com.fasterxml.jackson.core.JsonParser;
import com.fasterxml.jackson.core.JsonToken;
import io.vertx.core.json.JsonArray;
import io.vertx.core.json.JsonObject;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.util.function.Consumer;

/**
 * 解析数据流
 */
public class StreamingParser {

    public final static String ERROR_MSG = "Expected content to be an array";

    /**
     * 解析 JSON 数据流
     *
     * @param stream     JSON 数据流
     * @param batch_size 批量大小
     * @param consumer   数据消费者
     * @return 解析的记录数
     */
    public static int parseJson(InputStream stream, int batch_size, Consumer<JsonArray> consumer) throws IOException {
        int recordCount = 0;
        JsonFactory jsonfactory = new JsonFactory();
        try (JsonParser parser = jsonfactory.createParser(stream)) {
            if (parser.nextToken() != JsonToken.START_ARRAY)
                throw new IllegalStateException(ERROR_MSG);
            JsonArray records = new JsonArray();
            while (parser.nextToken() != JsonToken.END_ARRAY) {
                JsonObject obj = readJsonObject(parser);
                records.add(obj);
                if (records.size() == batch_size) {
                    recordCount += records.size();
                    consumer.accept(records);
                    records.clear();
                }
            }
            if (!records.isEmpty()) {
                recordCount += records.size();
                consumer.accept(records);
                records.clear();
            }
            return recordCount;
        }
    }

    /**
     * 解析 ES 导出的 JSON 数据文件
     *
     * @param stream     JSON 数据流
     * @param batch_size 批量大小
     * @param consumer   数据消费者
     * @return 解析的记录数
     */
    public static int parseESJson(InputStream stream, int batch_size, Consumer<JsonArray> consumer) throws IOException {
        int recordCount = 0;
        try (BufferedReader reader = new BufferedReader(new InputStreamReader(stream))) {
            JsonArray records = new JsonArray();
            do {
                String line = reader.readLine();
                if (line == null)
                    break;
                JsonObject obj = new JsonObject(line).getJsonObject("_source");
                if (obj == null)
                    continue;
                records.add(obj);
                if (records.size() == batch_size) {
                    recordCount += records.size();
                    consumer.accept(records);
                    records.clear();
                }
            } while (true);
            if (!records.isEmpty()) {
                recordCount += records.size();
                consumer.accept(records);
                records.clear();
            }
            return recordCount;
        }
    }

    /**
     * 读取对象信息
     *
     * @param parser JSON 解析器
     * @return JSON 对象
     */
    public static JsonObject readJsonObject(JsonParser parser) throws IOException {
        // Check the first token
        if (parser.currentToken() != JsonToken.START_OBJECT)
            throw new IllegalStateException("Expected content to be an object");
        JsonObject obj = new JsonObject();
        while (parser.nextToken() != JsonToken.END_OBJECT) {
            String key = parser.getCurrentName();
            JsonToken token = parser.nextToken();
            if (token == JsonToken.START_OBJECT)
                obj.put(key, readJsonObject(parser));
            else if (token == JsonToken.START_ARRAY)
                obj.put(key, readJsonArray(parser));
            if (token.isNumeric())
                obj.put(key, parser.getNumberValue());
            else if (token.isBoolean())
                obj.put(key, parser.getBooleanValue());
            else if (token == JsonToken.VALUE_STRING)
                obj.put(key, parser.getText());
        }
        return obj;
    }

    /**
     * 读取数组信息
     *
     * @param parser JSON 解析器
     * @return JSON 数组
     */
    public static JsonArray readJsonArray(JsonParser parser) throws IOException {
        // Check the first token
        if (parser.currentToken() != JsonToken.START_ARRAY)
            throw new IllegalStateException("Expected content to be an object");
        JsonArray array = new JsonArray();
        while (parser.nextToken() != JsonToken.END_ARRAY) {
            JsonToken token = parser.currentToken();
            if (token == JsonToken.START_OBJECT)
                array.add(readJsonObject(parser));
            if (token.isNumeric())
                array.add(parser.getNumberValue());
            else if (token.isBoolean())
                array.add(parser.getBooleanValue());
            else if (token == JsonToken.VALUE_STRING)
                array.add(parser.getText());
        }
        return array;
    }

}
